// In SecureFileTransportSystem/src/client/client_logic.c
// ... (includes and other functions)

#include <sys/stat.h> // For stat() to check if path is a directory
#include <limits.h>   // For PATH_MAX if needed for constructed path, though local_path should be large enough
#include <stdio.h>
#include <string.h>
#include <stdlib.h> // For free, calloc etc.
#include <errno.h>  // For errno and strerror()
#include <unistd.h> // For unlink()
#include <stdint.h> // For uintXX_t types
#include <sys/types.h>
#include <fcntl.h>

// Project-specific includes (adjust paths if necessary based on compiler errors)
#include "protocol.h"          // For message types, request/response structs
#include "crypto/sm4_crypto.h" // Assuming SM4_IV_SIZE, DOWNLOAD_BLOCK_PLAINTEXT_SIZE etc. are here
                               // Path might be different, e.g., "../include/crypto/sm4_crypto.h"
                               // or just "sm4_crypto.h" if include paths are set up broadly.

// Define a constant for max ciphertext block size expected from server (IV + data + padding)
#define CLIENT_DOWNLOAD_MAX_CIPHERTEXT_BLOCK (DOWNLOAD_BLOCK_PLAINTEXT_SIZE + SM4_IV_SIZE + SM4_BLOCK_SIZE)

// Forward declare get_next_sequence_number and get_error_message_string if not in a common client header
// uint16_t get_next_sequence_number();
// const char* get_error_message_string(uint8_t error_code);

int handle_download_file(int sockfd, const char *remote_path, const char *local_path,
                         const unsigned char *session_key, const unsigned char *session_token)
{
    if (!remote_path || !local_path || !session_key /* || !session_token */)
    { // Token optional for now in req struct
        fprintf(stderr, "[Download] Invalid arguments for download.\n");
        fflush(stderr);
        return -1;
    }

    printf("准备下载文件: %s -> %s\n", remote_path, local_path);
    fflush(stdout);

    // 1. 发送 MSG_TYPE_FILE_DOWNLOAD_REQUEST
    file_download_request_t down_req;
    strncpy(down_req.path, remote_path, sizeof(down_req.path) - 1);
    down_req.path[sizeof(down_req.path) - 1] = '\0';
    // memcpy(down_req.token, session_token, sizeof(down_req.token)); // If token was part of struct

    uint16_t req_seq = get_next_sequence_number(); // Assuming this function exists and is linked
    fprintf(stderr, "[CLIENT_DL_DEBUG] Sending FILE_DOWNLOAD_REQUEST, seq: %u, path: %s\n", req_seq, down_req.path);
    fflush(stderr);

    if (send_message(sockfd, MSG_TYPE_FILE_DOWNLOAD_REQUEST, req_seq, &down_req, sizeof(down_req)) != 0)
    {
        fprintf(stderr, "[Download] 发送下载请求失败.\n");
        fflush(stderr);
        return -1;
    }

    // 2. 接收 MSG_TYPE_FILE_DOWNLOAD_RESPONSE
    uint8_t resp_type;
    // uint16_t resp_seq; // recv_specific_message handles seq matching
    void *resp_data = NULL;
    uint32_t resp_len;

    fprintf(stderr, "[CLIENT_DL_DEBUG] Waiting for FILE_DOWNLOAD_RESPONSE with expected seq: %u\n", req_seq);
    fflush(stderr);

    int resp_ret = recv_specific_message(sockfd, MSG_TYPE_FILE_DOWNLOAD_RESPONSE, req_seq, &resp_data, &resp_len);
    if (resp_ret != 0)
    {
        fprintf(stderr, "[Download] 未能接收到下载响应或类型不匹配. resp_ret: %d\n", resp_ret);
        fflush(stderr);
        if (resp_data)
            free(resp_data);
        return -1;
    }

    fprintf(stderr, "[CLIENT_DL_DEBUG] Received FILE_DOWNLOAD_RESPONSE, len: %u\n", resp_len);
    fflush(stderr);

    if (!resp_data || resp_len < sizeof(file_download_response_t))
    {
        fprintf(stderr, "[Download] 下载响应数据无效.\n");
        fflush(stderr);
        if (resp_data)
            free(resp_data);
        return -1;
    }

    file_download_response_t *down_resp = (file_download_response_t *)resp_data;
    if (down_resp->result != ERR_SUCCESS)
    {
        fprintf(stderr, "[Download] 服务器拒绝下载请求: %s (错误码: %d)\n", get_error_message_string(down_resp->result), down_resp->result); // Assuming get_error_message_string exists
        fflush(stderr);
        free(resp_data);
        return -1;
    }

    uint32_t download_id = down_resp->download_id;
    uint64_t total_file_size = down_resp->size;
    // uint32_t server_plaintext_block_size = down_resp->block_size;
    fprintf(stderr, "[CLIENT_DL_DEBUG] 下载响应成功解析! ID: %u, 大小: %llu, 块大小: %u\n",
            download_id, (unsigned long long)total_file_size, down_resp->block_size);
    fflush(stderr);

    printf("[Download] 服务器接受下载. ID: %u, 文件大小: %llu bytes, 服务器块大小: %u bytes\n",
           download_id, (unsigned long long)total_file_size, down_resp->block_size);
    fflush(stdout);

    free(resp_data);
    resp_data = NULL;

    char final_local_path[PATH_MAX];
    strncpy(final_local_path, local_path, sizeof(final_local_path) - 1);
    final_local_path[sizeof(final_local_path) - 1] = '\0';

    fprintf(stderr, "[DEBUG_STDERR] Original local_path argument: '%s'\n", local_path);
    fprintf(stderr, "[DEBUG_STDERR] Initial final_local_path: '%s'\n", final_local_path);
    fflush(stderr);

    struct stat path_stat_debug;
    int stat_ret = stat(local_path, &path_stat_debug);

    if (stat_ret != 0)
    {
        fprintf(stderr, "[DEBUG_STDERR] stat(\"%s\") FAILED. errno: %d (%s)\n", local_path, errno, strerror(errno));
        fflush(stderr);
    }
    else
    {
        fprintf(stderr, "[DEBUG_STDERR] stat(\"%s\") SUCCEEDED.\n", local_path);
        fflush(stderr);
        if (S_ISDIR(path_stat_debug.st_mode))
        {
            fprintf(stderr, "[DEBUG_STDERR] S_ISDIR(\"%s\") is TRUE. Path is a directory.\n", local_path);
            fflush(stderr);

            const char *remote_filename = strrchr(remote_path, '/');
            if (remote_filename)
            {
                remote_filename++;
            }
            else
            {
                remote_filename = remote_path;
            }
            fprintf(stderr, "[DEBUG_STDERR] Determined remote_filename: '%s'\n", remote_filename);
            fflush(stderr);

            size_t dir_len = strlen(local_path);
            if (dir_len > 0 && local_path[dir_len - 1] != '/')
            {
                if ((dir_len + 1 + strlen(remote_filename) + 1) <= sizeof(final_local_path))
                {
                    snprintf(final_local_path, sizeof(final_local_path), "%s/%s", local_path, remote_filename);
                }
                else
                {
                    fprintf(stderr, "[Download] Path construction failed: buffer too small for %s/%s\n", local_path, remote_filename);
                    fflush(stderr);
                    return -1;
                }
            }
            else if (dir_len > 0 && local_path[dir_len - 1] == '/')
            {
                if ((dir_len + strlen(remote_filename) + 1) <= sizeof(final_local_path))
                {
                    snprintf(final_local_path, sizeof(final_local_path), "%s%s", local_path, remote_filename);
                }
                else
                {
                    fprintf(stderr, "[Download] Path construction failed: buffer too small for %s%s\n", local_path, remote_filename);
                    fflush(stderr);
                    return -1;
                }
            }
            else
            {
                if ((strlen(remote_filename) + 1) <= sizeof(final_local_path))
                {
                    strncpy(final_local_path, remote_filename, sizeof(final_local_path) - 1);
                    final_local_path[sizeof(final_local_path) - 1] = '\0';
                }
                else
                {
                    fprintf(stderr, "[Download] Path construction failed: buffer too small for %s (as filename)\n", remote_filename);
                    fflush(stderr);
                    return -1;
                }
            }
            fprintf(stderr, "[Download_STDERR] Local path was a directory. Effective save path: %s\n", final_local_path);
            fflush(stderr);
        }
        else
        {
            fprintf(stderr, "[DEBUG_STDERR] S_ISDIR(\"%s\") is FALSE. Path is not a directory.\n", local_path);
            fflush(stderr);
        }
    }

    FILE *local_fp = fopen(final_local_path, "wb");
    if (!local_fp)
    {
        fprintf(stderr, "[Download] 无法创建本地文件 '%s': %s (errno: %d)\n", final_local_path, strerror(errno), errno);
        fflush(stderr);
        return -1;
    }
    fprintf(stderr, "[Download_STDERR] Successfully opened local file for writing: '%s'\n", final_local_path);
    fflush(stderr);

    uint64_t received_file_size = 0;
    unsigned char download_iv[SM4_IV_SIZE];
    int iv_extracted = 0;
    uint32_t current_block_num = 0;
    int download_complete = 0;

    if (total_file_size == 0)
    {
        printf("[Download] 文件为空，下载完成.\n");
        fflush(stdout);
        download_complete = 1;
    }

    while (!download_complete)
    {
        file_download_data_request_t data_chunk_req;
        data_chunk_req.download_id = download_id;
        data_chunk_req.block_num = current_block_num;

        uint16_t data_req_seq = get_next_sequence_number();
        fprintf(stderr, "[CLIENT_DL_DEBUG] Loop iter: block_num=%u. Preparing to send DOWNLOAD_DATA_REQUEST (seq=%u)\n", current_block_num, data_req_seq);
        fflush(stderr);

        // 打印详细的请求内容
        fprintf(stderr, "[CLIENT_DL_DEBUG] DOWNLOAD_DATA_REQUEST details: ID: %u, Block: %u, Seq: %u\n",
                data_chunk_req.download_id, data_chunk_req.block_num, data_req_seq);
        fflush(stderr);

        if (send_message(sockfd, MSG_TYPE_FILE_DOWNLOAD_DATA_REQUEST, data_req_seq, &data_chunk_req, sizeof(data_chunk_req)) != 0)
        {
            fprintf(stderr, "[Download] 发送数据块 %u 请求失败.\n", current_block_num);
            fflush(stderr);
            fclose(local_fp);
            unlink(final_local_path);
            return -1;
        }
        fprintf(stderr, "[CLIENT_DL_DEBUG] Sent DOWNLOAD_DATA_REQUEST for block %u. Waiting for DOWNLOAD_DATA response...\n", current_block_num);
        fflush(stderr);

        resp_data = NULL;
        resp_len = 0;
        fprintf(stderr, "[CLIENT_DL_DEBUG] Calling recv_specific_message for DOWNLOAD_DATA with req_seq=%u\n", data_req_seq);
        fflush(stderr);

        // 添加重试逻辑
        int retry_count = 0;
        const int max_retries = 3;
        int recv_ret = -1;

        while (retry_count < max_retries)
        {
            recv_ret = recv_specific_message(sockfd, MSG_TYPE_FILE_DOWNLOAD_DATA, data_req_seq, &resp_data, &resp_len);

            if (recv_ret == 0)
            {
                // 成功接收到响应
                break;
            }
            else if (recv_ret == ERR_RECV_WOULD_BLOCK)
            {
                // 非阻塞socket暂时没有数据，等待后重试
                fprintf(stderr, "[CLIENT_DL_DEBUG] Received WOULD_BLOCK, will retry (%d/%d)...\n",
                        retry_count + 1, max_retries);
                fflush(stderr);

                // 使用一个简单延时循环替代usleep
                for (volatile int i = 0; i < 10000000; i++)
                {
                    // 空循环用于延时
                }

                retry_count++;
                continue;
            }
            else
            {
                // 其他错误，直接返回
                break;
            }
        }

        if (recv_ret != 0)
        {
            fprintf(stderr, "[Download] 未能接收到数据块 %u (seq %u) 或类型不匹配/超时. recv_ret: %d, errno: %d (%s), 重试次数: %d\n",
                    current_block_num, data_req_seq, recv_ret, errno, strerror(errno), retry_count);
            fflush(stderr);
            if (resp_data)
                free(resp_data);
            fclose(local_fp);
            unlink(final_local_path);
            return -1;
        }
        fprintf(stderr, "[CLIENT_DL_DEBUG] Received DOWNLOAD_DATA response for block %u. Resp_len: %u\n", current_block_num, resp_len);
        fflush(stderr);

        if (!resp_data || resp_len < offsetof(file_download_data_t, data))
        {
            fprintf(stderr, "[Download] 数据块 %u 响应数据无效 (太短 for header: %u bytes received).\n", current_block_num, resp_len);
            fflush(stderr);
            if (resp_data)
                free(resp_data);
            fclose(local_fp);
            unlink(final_local_path);
            return -1;
        }

        file_download_data_t *data_msg = (file_download_data_t *)resp_data;
        if (data_msg->download_id != download_id || data_msg->block_num != current_block_num)
        {
            fprintf(stderr, "[Download] 数据块 %u 响应ID或块号不匹配 (Expected ID %u, Blk %u; Got ID %u, Blk %u).\n",
                    current_block_num, download_id, current_block_num, data_msg->download_id, data_msg->block_num);
            fflush(stderr);
            if (resp_data)
                free(resp_data);
            fclose(local_fp);
            unlink(final_local_path);
            return -1;
        }
        fprintf(stderr, "[CLIENT_DL_DEBUG] Processing received block %u. Data_msg->data_len from server: %u\n", current_block_num, data_msg->data_len);
        fflush(stderr);

        const unsigned char *received_payload = data_msg->data;
        uint32_t received_payload_len = data_msg->data_len;

        unsigned char decrypted_buffer[DOWNLOAD_BLOCK_PLAINTEXT_SIZE + SM4_BLOCK_SIZE];
        size_t decrypted_data_len = 0;

        if (received_payload_len > 0)
        {
            if (current_block_num == 0 && !iv_extracted)
            {
                if (received_payload_len < SM4_IV_SIZE)
                {
                    fprintf(stderr, "[Download] Error: First data block is too short to contain IV (len: %u).\n", received_payload_len);
                    fflush(stderr);
                    if (resp_data)
                        free(resp_data);
                    fclose(local_fp);
                    unlink(final_local_path);
                    return -1;
                }
                memcpy(download_iv, received_payload, SM4_IV_SIZE);
                iv_extracted = 1;
                fprintf(stderr, "[CLIENT_DL_DEBUG] Extracted IV from block 0.\n");
                fflush(stderr);

                if (received_payload_len > SM4_IV_SIZE)
                {
                    decrypted_data_len = sizeof(decrypted_buffer);
                    if (sm4_cbc_decrypt_data(received_payload + SM4_IV_SIZE, received_payload_len - SM4_IV_SIZE,
                                             session_key, download_iv,
                                             decrypted_buffer, &decrypted_data_len) != 0)
                    {
                        fprintf(stderr, "[Download] Failed to decrypt first data block after IV.\n");
                        fflush(stderr);
                        if (resp_data)
                            free(resp_data);
                        fclose(local_fp);
                        unlink(final_local_path);
                        return -1;
                    }
                }
                else
                {
                    decrypted_data_len = 0;
                }
            }
            else if (iv_extracted)
            {
                decrypted_data_len = sizeof(decrypted_buffer);
                if (sm4_cbc_decrypt_data(received_payload, received_payload_len,
                                         session_key, download_iv,
                                         decrypted_buffer, &decrypted_data_len) != 0)
                {
                    fprintf(stderr, "[Download] Failed to decrypt data block %u.\n", current_block_num);
                    fflush(stderr);
                    if (resp_data)
                        free(resp_data);
                    fclose(local_fp);
                    unlink(final_local_path);
                    return -1;
                }
            }
            else
            {
                fprintf(stderr, "[Download] Error: IV not extracted for non-first block %u.\n", current_block_num);
                fflush(stderr);
                if (resp_data)
                    free(resp_data);
                fclose(local_fp);
                unlink(final_local_path);
                return -1;
            }
        }

        if (decrypted_data_len > 0)
        {
            if (fwrite(decrypted_buffer, 1, decrypted_data_len, local_fp) != decrypted_data_len)
            {
                fprintf(stderr, "[Download] Failed to write decrypted data to local file: %s\n", strerror(errno));
                fflush(stderr);
                if (resp_data)
                    free(resp_data);
                fclose(local_fp);
                unlink(final_local_path);
                return -1;
            }
        }
        received_file_size += decrypted_data_len;

        // fprintf(stdout, "[Download] Received block %u, plain_len: %zu. Total: %llu/%llu\n",
        //         current_block_num, decrypted_data_len, (unsigned long long)received_file_size, (unsigned long long)total_file_size);
        // fflush(stdout);

        int server_marked_last = data_msg->is_last;
        free(resp_data);
        resp_data = NULL;

        if (server_marked_last)
        {
            if (received_file_size != total_file_size && total_file_size > 0)
            {
                fprintf(stderr, "[Download] Warning: Total received size %llu does not match expected server size %llu.\n",
                        (unsigned long long)received_file_size, (unsigned long long)total_file_size);
                fflush(stderr);
            }
            download_complete = 1;
        }
        else if (total_file_size > 0 && received_file_size >= total_file_size)
        {
            download_complete = 1;
        }
        else if (total_file_size == 0 && server_marked_last)
        {
            download_complete = 1;
        }

        if (!download_complete)
        {
            current_block_num++;
        }
    }

    fclose(local_fp);
    local_fp = NULL;

    if (download_complete && received_file_size == total_file_size)
    {
        printf("文件下载成功: %s (%llu bytes)\n", final_local_path, (unsigned long long)received_file_size);
        fflush(stdout);
        return 0;
    }
    else
    {
        fprintf(stderr, "[Download] Download finished but with discrepancies or was incomplete. Received: %llu, Expected: %llu. Local file: %s\n",
                (unsigned long long)received_file_size, (unsigned long long)total_file_size, final_local_path);
        fflush(stderr);
        unlink(final_local_path);
        return -1;
    }
}

// Make sure to add a prototype for handle_download_file if it's called from main or other functions
// e.g. in a local header or at the top of client_logic.c
// int handle_download_file(int sockfd, const char *remote_path, const char *local_path,
//                        const unsigned char *session_key, const unsigned char *session_token);

// And update main or command parsing logic to call handle_download_file.
// Example in a command loop:
// else if (strncmp(command, "download ", 9) == 0) {
//     char remote_p[256], local_p[256];
//     if (sscanf(command + 9, "%s %s", remote_p, local_p) == 2) {
//         if (!g_auth_done) { printf("Please login first.\n"); continue; }
//         handle_download_file(sockfd, remote_p, local_p, g_session_key, g_session_token);
//     } else {
//         printf("Usage: download <remote_path> <local_path>\n");
//     }
// }